from dataset import MyDataset,DataCollator
import json
import torch
from modeling import fuseModelWithAppendToBegin
from transformers import AutoModel,AutoTokenizer,AutoModelForCausalLM

sent2vec_tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-large-zh-v1.5',model_max_length=512)
sent2vec = AutoModel.from_pretrained('BAAI/bge-large-zh-v1.5',torch_dtype=torch.float16)
llm_tokenizer = AutoTokenizer.from_pretrained('/data/lxy/baichuan2-chat-7b',trust_remote_code=True)
llm = AutoModelForCausalLM.from_pretrained('/data/lxy/baichuan2-chat-7b',trust_remote_code=True,torch_dtype=torch.bfloat16)
model=fuseModelWithAppendToBegin(sent2vec=sent2vec,sent2vec_tokenizer=sent2vec_tokenizer,
                                 llm=llm,llm_tokenizer=llm_tokenizer)
model.save_pretrained('/data/lxy/damnModel',max_shard='1GB')